#!/usr/bin/python2.4

# pylint: disable-msg=C6304

"""Framework for reading PPM files to extract drawn vectors."""

import os
import re
import subprocess
import sys
import urllib2

import matrix
from matrix import Point
from matrix import TransformStack
from PIL import Image

WORKDIR = os.path.join('.')

# (Red, Green, Blue)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)


def GetRed(color):
  return color[0]


def GetGreen(color):
  return color[1]


def GetBlue(color):
  return color[2]


class ImageReader(object):
  """Generic image reader.  Subclassed for UW/WRF WX and DFO current images."""

  def __init__(self, tesseract_api, trace_level=0):
    self._tesseract_api = tesseract_api
    self._trace_level = trace_level
    self._images = {}
    self._text_image_bytes = {}

  def ReadImages(self, desired, output_file=sys.stdout):
    """Read all of the images determined by GetSourceFiles.

    Args:
      desired: Set of standard names describing desired images,
        typically 'today', or UW model init times in UTC
        formatted '%Y%m%d%H'.
      output_file: Where to write the CSV data extracted from the images.
    """
    for file_name in self.GetSourceFiles(desired):
      try:
        self._images[file_name], self._text_image_bytes[file_name] = (
            self.ReadImage(file_name))
        self.CustomAnalysis(file_name, output_file)
      except IOError, e:
        print '%s' % (str(e),)
      output_file.flush()

    if not self._images:
      print 'No images loaded?'
      return
    self._file_names = self._images.keys()
    self._file_names.sort()
    self._canonical_file_name = self._file_names[0]
    if False:
      self.CustomAnalysisAll()

  def ReadImage(self, file_name):
    """Reads image from file_name, and crops it per bounding box.

    Selects out the colors of interest.
    Runs the text conversion on the image to extract custom fields
    holding the date and time the model is valid.

    Args:
      file_name: The name of the file holding the PPM data to read.

    Returns:
      A pair: the image, and the label (str) extracted from the image.
    """
    if True or self._trace_level >= 1:
      print 'opening %s' % (file_name,)
    image = Image.open(file_name)

    if self._trace_level >= 1:
      print 'size=%d %d' % (image.size)

    label = self.ConvertText(image, file_name)

    crop_box = self.ImageBoundingBox()

    image = image.crop(crop_box)

    image = self.SelectPreferredColor(file_name, image)
    return image, label

  def ConvertText(self, image, file_name):
    """Extracts the text part of the image, and runs through the Tesseract OCR.

    Sets self._unix_epoch_time as a side effect.

    Args:
      image: The image to analyze.
      file_name: The name of the file holding the original image.

    Returns:
      The cleaned up OCR string.
    """
    if self._tesseract_api:
      image_text = image.crop(text_box)
      text_box = self.TextBoundingBox()
      # Save just the text in its own ppm file, which can be given
      # directly to Tesseract.
      # TODO(rrh): We can only do this if we read a ppm file; issues
      # with getting palettes for GIF files?
      output_fname = os.path.join(
          os.path.dirname(file_name),
          'text.' + os.path.basename(file_name))
      image_text.save(output_fname, 'PPM')
      text_image_bytes = open(output_fname, 'r').read()
      string_scanned = self._tesseract_api.Recognize(text_image_bytes)
    else:
      string_scanned = '?'

    string_scanned = string_scanned.replace('\n', ' ')
    return self.GetImageTime(string_scanned, file_name)


def SelectPureRed(base_image):
  """Selects out pure red pixels.  Used to find windbarbs on UW wind images.

  Args:
    base_image: The base image to discard non-red points.

  Returns:
    base_image modified to retain just interesting red pixels.
  """
  red_data = []
  for p0 in base_image.getdata():
    if True:
      # Carry image through directly
      red_data.append(p0)
    else:
      if p0 == RED:
        red_data.append(RED)
      else:
        red_data.append(WHITE)
  base_image.putdata(red_data)

  if False:
    # Make marks to help understand the coordinate systems.
    # TODO(rrh): These numbers are specific to UW/WRF images.
    # X increases from left to right
    # Y increases from top to bottom (X11 coordinate system)
    for y in [10, 11, 12, 13, 14, 15, 16]:
      for x in [10, 11, 12, 13, 14]:
        base_image.putpixel((x, y), GREEN)  # upper left corner
        # base_image.putpixel((628-x, 628-y), BLUE)  # lower right corner

  if False:
    # Draw a wind barb to experiment with parameters of the barb.
    s = TransformStack()
    s.Translate(Point(45, 104))
    s.Scale(1, -1)
    s.Scale(12, 12)

    class WrappedImage(object):
      def __init__(self, base_image):
        self._base_image = base_image

      def SetPixel(self, point, trace=False):  # pylint: disable-msg=W0613
        base_image.putpixel((point.x, point.y), BLACK)

    wi = WrappedImage(base_image)
    true_wind_speed = 15
    true_wind_angle = 205
    matrix.RasterWindBarbGenerator(
        wi, s,
        true_wind_speed, true_wind_angle,  # TODO(rrh): convert to using Polar
        color=RED,
        rotation=-60, barbette_length=0.5, barbette_spacing=0.30)

  return base_image


def AndRed(base_image, rhs_image):
  """Combines two images, preserving pixels that are red in both; else white.

  Args:
    base_image: The starting image; will be modified.
    rhs_image: Another image to "AND" onto base_image.

  Returns:
    base_image modified to represent the "AND" of red pixels.
  """

  and_data = []
  for p0, p1 in zip(base_image.getdata(), rhs_image.getdata()):
    if p0 == RED and p1 == RED:
      and_data.append(RED)
    else:
      and_data.append(WHITE)
  base_image.putdata(and_data)
  return base_image


def MajorityRed(rhs_images):
  """Does color separation looking for Red pixels exceeding 75% threshhold."""
  counts = {}
  for image_key in rhs_images:
    print >>sys.stdout, 'Majority-combine in %s' % (image_key,)
    image = rhs_images[image_key]
    pixel_number = 0
    for pixel in image.getdata():
      counts.setdefault(pixel_number, 0)
      if pixel == RED:
        counts[pixel_number] += 1
      pixel_number += 1

  maxcount = len(rhs_images)
  new_image = rhs_images[rhs_images.keys()[0]].copy()
  new_pixels = []

  pixel_number = 0
  for pixel in image.getdata():
    if counts[pixel_number] > 0.75 * maxcount:
      new_pixels.append(RED)
    else:
      new_pixels.append(WHITE)
    pixel_number += 1
  new_image.putdata(new_pixels)
  return new_image


def OrRed(base_image, rhs_image):
  or_data = []
  for p0, p1 in zip(base_image.getdata(), rhs_image.getdata()):
    if p0 == RED or p1 == RED:
      or_data.append(RED)
    else:
      or_data.append(WHITE)
  base_image.putdata(or_data)
  return base_image


def UnTesseract(tesseract_output_str):
  """Replace characters to undo effects of bad Tesseract conversion."""
  return (
      tesseract_output_str
      .replace('B', '8')
      .replace('U', '0')
      .replace('D', '0')
      .replace('O', '0')
      .replace('I', '1')
      )


def GetGIFFromServer(http_dir, local_dir, base_fmt, last_delta_hour):
  """Read a .gif file from some web server, save and convert to a ppm file.

  Args:
    http_dir: Full URL to the directory holding the files.
    local_dir: Full path to the local file system's directory.
    base_fmt: A format string used to construct the names of gif files
      held in the result directory.
    last_delta_hour: How many hour steps to get .gif files for.

  Returns:
    A pair: (success_count, fail_count).
  """
  success_count = 0
  fail_count = 0

  print 'mkdir -p %s' % (local_dir,)
  try:
    os.makedirs(local_dir)
  except OSError, unused_e:
    pass

  for hour in xrange(1, last_delta_hour+1):
    base_name = base_fmt % (hour,)
    full_url = http_dir + base_name
    try:
      url = urllib2.urlopen(http_dir + '/' + base_name)
      contents = url.read()
      url.close()
      gif_file_name = os.path.join(local_dir, base_name)
      print 'read from %s total of %d bytes write to %s' % (
          full_url, len(contents), gif_file_name)
      outfd = open(gif_file_name, 'w')
      outfd.write(contents)
      outfd.close()

      if True:
        # TODO(rrh): Figure out why we can't read and process GIF files.
        ConvertGIFToPPM(gif_file_name)
      success_count += 1

    except urllib2.HTTPError, unused_e:
      print '%s: Not present on server' % (full_url)
      fail_count += 1
  return success_count, fail_count


def ConvertGIFToPPM(gif_file_name):
  """Convert the gif held in gif_file_name into a sibling .ppm file."""
  ppm_file_name = re.sub(r'\.gif$', '.ppm', gif_file_name)
  # print 'ppm_file_name=%s' % (ppm_file_name,)

  if False:
    # TODO(rrh): I couldn't get the Python Imaging Library to write a PPM
    # given a GIF.
    # There were issues with IOError: cannot write mode P as PPM,
    # which presumably involve palette choices.
    # Figure this out, as otherwise we need a subprocess giftopnm,
    # which is not very portable.
    image = Image.open(gif_file_name)
    image.load()
    print 'format=%s size=(%d, %d), mode=%s' % (
        image.format, image.size[0], image.size[1], image.mode)
    image.save(ppm_file_name, 'PPM')
  else:
    output_fd = open(ppm_file_name, 'w')
    p = subprocess.Popen(['giftopnm', gif_file_name], stdout=output_fd)
    p.wait()
    output_fd.close()
  return ppm_file_name
